rm(list = ls())
library(tidyverse)
library(lubridate)

if(Sys.info()["user"]=="XXXXX") setwd("XXXXX")
if(Sys.info()["user"]=="devap") setwd("C:/Users/devap/Dropbox/Vaccines/SMS_ADMIN/")
if(Sys.info()["user"]=="jimmy") setwd("~/Dropbox/SMS_ADMIN/")

#D vector based on shares
ProportionalAssignment = function(Shares, Nt) {
  k = length(Shares)
  n_floor = floor(Nt * Shares) #number of units assigned to each treatment, rounded down
  remainder = Nt * Shares - n_floor
  units_remaining = Nt - sum(n_floor) #remaining number of units to be assigned
  
  if (units_remaining > 0) {
    Dt = c(rep(1:k, n_floor),
           # assigning each treatment acoording to rounded down average count
           sample(
             1:k,
             size = units_remaining,
             replace = T,
             prob = remainder
           )) #remaining units assigned randomly from remain replicate assignments
  } else {
    Dt = rep(1:k, n_floor)
  }
  sample(Dt)
}


DtchoiceThompson = function(Y, D, #outcomes and treatments thus far
                            k, #number of treatments
                            C, #vector of treatment cost
                            Nt) {
  # number of observations for period t
  
  A = 1 + tapply(Y, D, sum, default = 0) 
  B = 1 + tapply(1-Y, D, sum, default = 0) 
  
  Dt = rep(0, Nt)
  previousD = -Inf # auxiliary variable to avoid repeat assignments of same D
  
  for (i  in 1:Nt) {
    thetadraw = sapply(1:k, function(j)
      rbeta(1, A[j], B[j]))
    Dt[i] = which.max(thetadraw - C)
    previousD = Dt[i]
  }
  
  factor(Dt, levels = 1:k)
}



DtchoiceThompsonProbabilities = function(Y, D, #outcomes and treatments thus far
                                         k, #number of treatments
                                         C = rep(0, k), #vector of treatment cost
                                         RR = 50000) {
  #number of replication draws
  
  # Repeat Thompson sampling RR times
  DtRR = DtchoiceThompson(Y, D, k, C, RR)
  P_Dt = table(DtRR) / RR #average count for each treatment value and covariate value, replicated sample
  
  P_Dt = as_tibble(matrix(P_Dt, 1, k))
  colnames(P_Dt) = paste(1:k)
  P_Dt
}


DtchoiceThompson_modified = function(alpha) {
  if (max(alpha) < 1) {
    Shares = alpha * (1 - alpha) # Based on stationary distribution for modified Thompson
    Shares = Shares / sum(Shares)
  } else {
    Shares = alpha
  }
  
  return(Shares)
}

#priorname = paste("OUTPUT/COVID/Priors/", Sys.Date(), "priors.csv", sep = "")
priorname = paste("OUTPUT/COVID/Priors/2020-10-07priors.csv", sep = "")
priordata= read_csv(priorname)  %>% 
    mutate(treatment=factor(treatment))

by_treatment= priordata %>% 
    group_by(treatment, .drop=FALSE) %>% 
    summarise(avg=mean(outcome), count=n(), successes=as.integer(sum(outcome)))
alpha = DtchoiceThompsonProbabilities(
    priordata$outcome,
    priordata$treatment,
    k=length(levels(priordata$treatment)),
  ) %>% as_vector

P_current = DtchoiceThompson_modified(alpha)
P_current_tibble =  tibble(share = as_vector(P_current), treatment = factor(levels(priordata$treatment)))
treatment_assignment = tibble(treatment=ProportionalAssignment(P_current, 100000))
filename = paste("OUTPUT/COVID/Treatment_assignments/", Sys.Date(), "_t_assignments.csv", sep = "")
write_csv(treatment_assignment, filename)


Bayes_table=by_treatment %>%
    mutate(alpha_d= 1 + successes,
           beta= 1 + count-successes,
           mean = alpha_d / (alpha_d+beta),
           var = alpha_d * beta / ((alpha_d+beta)^2 * (alpha_d+beta +1)  ),
           std = sqrt(var),
           succ_proba = alpha) %>%
    select(treatment, mean, std, succ_proba)
  
  colnames(Bayes_table)=c("Treatment", "Mean", "Standard dev", "Probability optimal")
  
  Bayes_table
  filename = paste("OUTPUT/COVID/Bayes_tables/", Sys.Date(), "_bayes_table.csv", sep = "")
  write_csv(Bayes_table, filename)

